#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
新版本特征提取脚本 - 适配新的文件结构

主要特点：
1. 直接从sourceData和targetDataset文件夹读取mat文件
2. 基于task1的特征提取逻辑
3. 生成源域和目标域的特征CSV文件
4. 生成分析报告和可视化图表
"""

import os
import sys
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# 添加当前目录到路径
sys.path.append('.')

# 导入自定义模块
from feature_extractor import FeatureExtractor
from feature_selector import FeatureSelector
from data_analyzer import DataAnalyzer

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class NewDataLoader:
    """新的数据加载器 - 适配新的文件结构"""
    
    def __init__(self, source_data_path, target_data_path, target_sampling_rate=12000):
        self.source_data_path = source_data_path
        self.target_data_path = target_data_path
        self.target_fs = target_sampling_rate
        
        # 轴承参数信息
        self.bearing_params = {
            'SKF6205': {'n': 9, 'd': 0.3126, 'D': 1.537},  # 驱动端轴承
            'SKF6203': {'n': 9, 'd': 0.2656, 'D': 1.122}   # 风扇端轴承
        }
    
    def load_source_data(self):
        """加载源域数据"""
        print("📂 开始加载源域数据...")
        
        source_files = []
        for file_name in os.listdir(self.source_data_path):
            if file_name.endswith('.mat'):
                source_files.append(os.path.join(self.source_data_path, file_name))
        
        print(f"发现源域文件数量: {len(source_files)}")
        
        # 加载源域文件
        source_data = self._load_mat_files(source_files, 'source')
        
        return source_data
    
    def load_target_data(self):
        """加载目标域数据"""
        print("📂 开始加载目标域数据...")
        
        target_files = []
        for file_name in os.listdir(self.target_data_path):
            if file_name.endswith('.mat'):
                target_files.append(os.path.join(self.target_data_path, file_name))
        
        print(f"发现目标域文件数量: {len(target_files)}")
        
        # 加载目标域文件
        target_data = self._load_mat_files(target_files, 'target')
        
        return target_data
    
    def _load_mat_files(self, file_list, data_type):
        """加载MAT文件数据"""
        import scipy.io as sio
        from tqdm import tqdm
        
        all_data = []
        
        for file_path in tqdm(file_list, desc=f"Loading {data_type} files"):
            try:
                data = sio.loadmat(file_path)
                file_name = os.path.basename(file_path)
                
                # 提取信号数据
                de_signal = self._extract_signal(data, 'DE')
                fe_signal = self._extract_signal(data, 'FE')
                ba_signal = self._extract_signal(data, 'BA')
                rpm = self._extract_rpm(data)
                
                # 提取文件信息
                file_info = self._parse_file_info(file_name, data_type)
                
                all_data.append({
                    'file_name': file_name,
                    'de_signal': de_signal,
                    'fe_signal': fe_signal,
                    'ba_signal': ba_signal,
                    'rpm': rpm,
                    'sampling_rate': file_info['sampling_rate'],
                    'fault_type': file_info['fault_type'],
                    'fault_size': file_info['fault_size'],
                    'load_condition': file_info['load_condition'],
                    'data_type': data_type
                })
                
            except Exception as e:
                print(f"⚠️ 加载文件失败 {file_path}: {e}")
                continue
        
        print(f"✅ 成功加载 {len(all_data)} 个{data_type}文件")
        return all_data
    
    def _extract_signal(self, data, signal_type):
        """提取信号数据"""
        # 查找包含信号类型的键
        for key in data.keys():
            if signal_type in key and 'time' in key:
                signal_data = data[key].flatten()
                # 重采样到目标采样率
                if len(signal_data) > 0:
                    return self._resample_signal(signal_data, signal_type)
        
        # 对于目标域数据，尝试直接使用单字母键
        if signal_type == 'DE':
            for key in data.keys():
                if len(key) == 1 and key.isalpha():  # 单字母键如'A', 'B'等
                    signal_data = data[key].flatten()
                    if len(signal_data) > 0:
                        return self._resample_signal(signal_data, signal_type)
        
        # 如果没找到，返回空数组
        return np.array([])
    
    def _resample_signal(self, signal, signal_type):
        """重采样信号到目标采样率"""
        if len(signal) == 0:
            return signal
            
        # 根据信号类型确定原始采样率
        if signal_type == 'DE':
            original_fs = 12000  # 驱动端12kHz
        else:
            original_fs = 12000  # 风扇端和基座12kHz
        
        if original_fs == self.target_fs:
            return signal
        
        # 计算重采样参数
        import scipy.signal as signal_lib
        num_samples = int(len(signal) * self.target_fs / original_fs)
        resampled_signal = signal_lib.resample(signal, num_samples)
        
        return resampled_signal
    
    def _extract_rpm(self, data):
        """提取RPM数据"""
        for key in data.keys():
            if 'rpm' in key.lower():
                return float(data[key][0][0])
        return 1796.0  # 默认RPM
    
    def _parse_file_info(self, file_name, data_type):
        """解析文件信息"""
        if data_type == 'target':
            return {
                'sampling_rate': 32000,  # 目标域32kHz
                'fault_type': 'Unknown',
                'fault_size': 0.0,
                'load_condition': 0
            }
        
        # 解析源域文件信息
        sampling_rate = 12000
        if '48kHz' in file_name:
            sampling_rate = 48000
        
        # 解析故障类型
        fault_type = 'Normal'
        if '-B-' in file_name:
            fault_type = 'Ball'
        elif '-IR-' in file_name:
            fault_type = 'Inner Race'
        elif '-OR-' in file_name:
            fault_type = 'Outer Race'
        elif 'Normal_data' in file_name or file_name.startswith('N_'):
            fault_type = 'Normal'
        
        # 解析故障尺寸
        fault_size = 0.0
        for size in ['0007', '0014', '0021', '0028']:
            if size in file_name:
                fault_size = float(size) / 1000
                break
        
        # 解析载荷条件
        load_condition = 0
        for load in ['0', '1', '2', '3']:
            if f'_{load}.mat' in file_name:
                load_condition = int(load)
                break
        
        return {
            'sampling_rate': sampling_rate,
            'fault_type': fault_type,
            'fault_size': fault_size,
            'load_condition': load_condition
        }


def extract_features_from_data(data_list, feature_extractor, data_type):
    """从数据列表中提取特征"""
    features_list = []
    
    for data in data_list:
        try:
            # 使用DE信号进行特征提取
            de_signal = data['de_signal']
            if len(de_signal) > 0:
                features = feature_extractor.extract_all_features(
                    de_signal, 
                    data['rpm'], 
                    'SKF6205'
                )
                features['file_name'] = data['file_name']
                features['fault_type'] = data['fault_type']
                features['fault_size'] = data['fault_size']
                features['load_condition'] = data['load_condition']
                features['data_type'] = data_type
                features_list.append(features)
        except Exception as e:
            print(f"⚠️ 特征提取失败 {data['file_name']}: {e}")
            continue
    
    return features_list


def generate_feature_analysis_visualizations(feature_selector, X, y, importance_df, corr_matrix, 
                                           distribution_stats, output_dir, timestamp):
    """生成特征分析可视化"""
    print("📊 生成特征分析可视化...")
    
    # 1. 特征重要性图
    importance_path = f"{output_dir}/feature_importance_analysis_{timestamp}.png"
    feature_selector.plot_feature_importance(importance_df, top_n=20, save_path=importance_path)
    
    # 2. 特征相关性热力图
    correlation_path = f"{output_dir}/feature_correlation_heatmap_{timestamp}.png"
    feature_selector.plot_feature_correlation(corr_matrix, save_path=correlation_path)
    
    # 3. 特征分布图
    distribution_path = f"{output_dir}/feature_distribution_{timestamp}.png"
    feature_selector.plot_feature_distribution(X, y, X.columns.tolist(), save_path=distribution_path)
    
    # 4. 特征空间可视化
    space_path = f"{output_dir}/feature_space_visualization_{timestamp}.png"
    feature_selector.visualize_features(X, y, method='tsne', save_path=space_path)
    
    print("✅ 特征分析可视化完成")


def generate_comprehensive_report(source_data, target_data, importance_df, corr_matrix, high_corr_pairs,
                                distribution_stats, selected_features, output_dir, timestamp):
    """生成综合报告"""
    print("📝 生成综合报告...")
    
    report = f"""# 新版本特征提取 - 综合报告

## 1. 任务概述

### 1.1 任务目标
基于新的文件结构，完成以下任务：
1. 从sourceData文件夹中提取源域数据特征
2. 从targetDataset文件夹中提取目标域数据特征
3. 进行特征分析和选择
4. 生成源域和目标域的特征CSV文件

### 1.2 数据概况
- **源域数据**: {len(source_data)}个文件
- **目标域数据**: {len(target_data)}个文件
- **特征维度**: {len(selected_features)}个选中特征

## 2. 数据来源

### 2.1 源域数据
源域数据来自sourceData文件夹，包含以下类型：
- 12kHz_DE_data: 12kHz采样频率的驱动端数据
- 12kHz_FE_data: 12kHz采样频率的风扇端数据
- 48kHz_DE_data: 48kHz采样频率的驱动端数据
- 48kHz_Normal_data: 48kHz采样频率的正常数据

### 2.2 目标域数据
目标域数据来自targetDataset文件夹，包含A-P共16个未知标签的故障文件。

## 3. 特征提取

### 3.1 特征类型
提取了以下5类特征：
1. **时域特征**: 均值、标准差、RMS、偏度、峰度等
2. **频域特征**: 频域统计、频带能量等
3. **时频域特征**: 小波变换、STFT特征等
4. **故障特征频率**: 轴承故障特征频率及其谐波
5. **非线性特征**: 近似熵、样本熵、排列熵等

### 3.2 特征选择结果
使用随机森林方法选择了{len(selected_features)}个重要特征：

"""
    
    # 添加选中的特征
    for i, feature in enumerate(selected_features, 1):
        report += f"{i}. {feature}\n"
    
    report += f"""
## 4. 特征重要性分析

### 4.1 前10个重要特征
"""
    
    # 添加前10个重要特征
    top_10_features = importance_df.head(10)
    for i, (_, row) in enumerate(top_10_features.iterrows(), 1):
        report += f"{i}. **{row['feature']}**: 综合分数 {row['combined_score']:.4f}\n"
    
    report += f"""
### 4.2 特征重要性统计
- **随机森林平均重要性**: {importance_df['random_forest'].mean():.4f}
- **互信息平均分数**: {importance_df['mutual_info'].mean():.4f}
- **F检验平均分数**: {importance_df['f_test'].mean():.4f}
- **综合平均分数**: {importance_df['combined_score'].mean():.4f}

## 5. 数据保存

### 5.1 保存文件
1. **源域特征数据**: `source_domain_features_{timestamp}.csv`
2. **目标域特征数据**: `target_domain_features_{timestamp}.csv`
3. **源域选中特征**: `source_domain_selected_features_{timestamp}.csv`
4. **目标域选中特征**: `target_domain_selected_features_{timestamp}.csv`

### 5.2 可视化文件
1. **特征重要性图**: `feature_importance_analysis_{timestamp}.png`
2. **特征相关性图**: `feature_correlation_heatmap_{timestamp}.png`
3. **特征分布图**: `feature_distribution_{timestamp}.png`
4. **特征空间图**: `feature_space_visualization_{timestamp}.png`

## 6. 技术实现

### 6.1 数据加载
- 直接从sourceData和targetDataset文件夹读取mat文件
- 自动解析文件信息（故障类型、故障尺寸、载荷条件等）
- 统一重采样到12kHz采样频率

### 6.2 特征提取方法
- 基于task1的特征提取逻辑
- 支持时域、频域、时频域、故障频率、非线性特征提取
- 针对轴承故障特点优化特征提取算法

### 6.3 特征选择策略
- 基于随机森林、互信息、F检验的综合评分
- 考虑特征相关性和分布差异进行选择
- 选择最具区分性的特征用于后续分析

## 7. 结论与建议

### 7.1 主要成果
1. 成功从新的文件结构中提取了源域和目标域特征
2. 完成了特征分析和选择，选择了{len(selected_features)}个重要特征
3. 生成了源域和目标域的特征CSV文件，为后续分析奠定了基础
4. 提供了完整的可视化分析和报告

### 7.2 技术特点
1. **适配性强**: 能够处理新的文件结构
2. **特征全面**: 提取5类46维特征，全面描述故障信息
3. **科学选择**: 基于多种方法综合评估特征重要性
4. **完整输出**: 提供CSV文件、可视化图表和详细报告

### 7.3 后续建议
1. 使用生成的特征CSV文件进行后续的机器学习任务
2. 基于特征重要性进行进一步的特征工程
3. 利用源域和目标域的特征差异进行迁移学习研究

---
*报告生成时间: {datetime.now().strftime('%Y年%m月%d日 %H:%M:%S')}*
*基于新文件结构的数据分析结果生成*
"""
    
    # 保存报告
    report_path = f"{output_dir}/feature_extraction_report_{timestamp}.md"
    with open(report_path, 'w', encoding='utf-8') as f:
        f.write(report)
    
    print(f"✅ 综合报告已保存: {report_path}")


def main():
    """主函数"""
    print("🚀 开始执行新版本特征提取任务")
    print("=" * 80)
    
    # 记录开始时间
    start_time = datetime.now()
    timestamp = start_time.strftime("%Y%m%d_%H%M%S")
    print(f"⏰ 执行时间: {timestamp}")
    
    # 创建输出目录
    output_dir = '.'
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # 步骤1: 数据加载
        print("\n" + "="*60)
        print("步骤1: 数据加载")
        print("="*60)
        
        data_loader = NewDataLoader(
            source_data_path='../datasets/sourceData',
            target_data_path='../datasets/targetDataset',
            target_sampling_rate=12000
        )
        
        # 加载源域数据
        source_data = data_loader.load_source_data()
        
        # 加载目标域数据
        target_data = data_loader.load_target_data()
        
        print(f"✅ 数据加载完成:")
        print(f"   源域数据: {len(source_data)}个文件")
        print(f"   目标域数据: {len(target_data)}个文件")
        
        # 步骤2: 特征提取
        print("\n" + "="*60)
        print("步骤2: 特征提取")
        print("="*60)
        
        # 创建特征提取器
        feature_extractor = FeatureExtractor(target_fs=12000)
        
        # 提取源域特征
        print("🔍 提取源域特征...")
        source_features = extract_features_from_data(source_data, feature_extractor, 'source')
        
        # 提取目标域特征
        print("🔍 提取目标域特征...")
        target_features = extract_features_from_data(target_data, feature_extractor, 'target')
        
        print(f"✅ 特征提取完成:")
        print(f"   源域特征: {len(source_features)}个样本")
        print(f"   目标域特征: {len(target_features)}个样本")
        
        # 步骤3: 特征选择
        print("\n" + "="*60)
        print("步骤3: 特征选择和可视化")
        print("="*60)
        
        # 准备特征选择数据
        source_df = pd.DataFrame(source_features)
        target_df = pd.DataFrame(target_features)
        
        # 获取数值特征列
        numeric_columns = source_df.select_dtypes(include=[np.number]).columns
        numeric_columns = [col for col in numeric_columns if col not in ['fault_size', 'load_condition']]
        
        X_source = source_df[numeric_columns]
        y_source = source_df['fault_type']
        
        # 创建特征选择器
        feature_selector = FeatureSelector()
        
        # 分析特征重要性
        importance_df = feature_selector.analyze_feature_importance(X_source, y_source)
        
        # 分析特征相关性
        corr_matrix, high_corr_pairs = feature_selector.analyze_feature_correlation(X_source)
        
        # 分析特征分布
        distribution_stats = feature_selector.analyze_feature_distribution(X_source, y_source)
        
        # 选择特征
        X_selected = feature_selector.select_features(X_source, y_source, method='random_forest', k=30)
        
        # 生成特征分析可视化
        generate_feature_analysis_visualizations(
            feature_selector, X_source, y_source, importance_df, corr_matrix, 
            distribution_stats, output_dir, timestamp
        )
        
        # 步骤4: 数据保存
        print("\n" + "="*60)
        print("步骤4: 数据保存")
        print("="*60)
        
        # 保存完整的特征数据
        source_file = f"{output_dir}/source_domain_features_{timestamp}.csv"
        target_file = f"{output_dir}/target_domain_features_{timestamp}.csv"
        
        source_df.to_csv(source_file, index=False, encoding='utf-8-sig')
        target_df.to_csv(target_file, index=False, encoding='utf-8-sig')
        
        # 保存选中的特征数据
        source_selected_df = source_df[['file_name', 'fault_type', 'fault_size', 'load_condition'] + feature_selector.selected_features]
        target_selected_df = target_df[['file_name'] + feature_selector.selected_features]
        
        source_selected_file = f"{output_dir}/source_domain_selected_features_{timestamp}.csv"
        target_selected_file = f"{output_dir}/target_domain_selected_features_{timestamp}.csv"
        
        source_selected_df.to_csv(source_selected_file, index=False, encoding='utf-8-sig')
        target_selected_df.to_csv(target_selected_file, index=False, encoding='utf-8-sig')
        
        print(f"✅ 数据保存完成:")
        print(f"   源域完整特征: {source_file}")
        print(f"   目标域完整特征: {target_file}")
        print(f"   源域选中特征: {source_selected_file}")
        print(f"   目标域选中特征: {target_selected_file}")
        
        # 步骤5: 生成综合报告
        print("\n" + "="*60)
        print("步骤5: 生成综合报告")
        print("="*60)
        
        generate_comprehensive_report(
            source_data, target_data, importance_df, corr_matrix, high_corr_pairs,
            distribution_stats, feature_selector.selected_features,
            output_dir, timestamp
        )
        
        # 计算执行时间
        end_time = datetime.now()
        execution_time = end_time - start_time
        
        print("\n" + "="*80)
        print("🎉 新版本特征提取任务执行完成！")
        print("="*80)
        print(f"⏱️ 总执行时间: {execution_time}")
        print(f"📁 输出目录: {output_dir}")
        print("="*80)
        
    except Exception as e:
        print(f"❌ 执行过程中出现错误: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    return True


if __name__ == "__main__":
    success = main()
    if success:
        print("\n🎉 新版本特征提取任务执行成功！")
    else:
        print("\n❌ 新版本特征提取任务执行失败！")
